import requests
import json
import time


class MCPClient:
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url
        self.ctx_id = None

    def chat_stream(self, message: str):
        """流式聊天接口"""
        url = f"{self.base_url}/chat"
        payload = {"message": message}
        if self.ctx_id:
            payload["ctx_id"] = self.ctx_id

        response = requests.post(
            url,
            json=payload,
            stream=True,
            headers={"Accept": "text/event-stream"}
        )

        buffer = ""
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                buffer += chunk.decode('utf-8')
                while "\n\n" in buffer:
                    event_data, buffer = buffer.split("\n\n", 1)
                    event_lines = event_data.strip().split('\n')

                    event_type = None
                    data = None

                    for line in event_lines:
                        if line.startswith('event:'):
                            event_type = line.split(':', 1)[1].strip()
                        elif line.startswith('data:'):
                            data = json.loads(line.split(':', 1)[1].strip())

                    if event_type and data:
                        # 处理上下文ID
                        if 'ctx_id' in data and not self.ctx_id:
                            self.ctx_id = data['ctx_id']

                        # 根据事件类型处理
                        if event_type == "token":
                            yield data['token']
                        elif event_type == "context_update":
                            print(f"\n[Context Updated] Tokens: {data['token_count']}")
                        elif event_type == "done":
                            print("\n[Generation Completed]")


# 使用示例
if __name__ == "__main__":
    client = MCPClient()

    while True:
        message = input("\nYou: ")
        if message.lower() in ["exit", "quit"]:
            break

        print("AI: ", end="", flush=True)
        for token in client.chat_stream(message):
            print(token, end="", flush=True)

        print(f"\nCurrent Context ID: {client.ctx_id}")